from collections import defaultdict
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import scipy
import scipy.io
from sklearn.preprocessing import label_binarize
import torch_geometric.transforms as T
from data_utils import even_quantile_labels, to_sparse_tensor

from torch_geometric.datasets import Planetoid, Amazon, Coauthor, Twitch, PPI, Reddit
from torch_geometric.transforms import NormalizeFeatures, RadiusGraph
from torch_geometric.data import Data, Batch
from torch_geometric.utils import stochastic_blockmodel_graph, subgraph, homophily, to_dense_adj, dense_to_sparse

from torch_geometric.nn import GCNConv, SGConv, SAGEConv, GATConv


import pickle as pkl
import os

class GCN_gen(nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels, num_layers):
        super(GCN_gen, self).__init__()

        self.convs = nn.ModuleList()
        self.convs.append(
            GCNConv(in_channels, hidden_channels))
        for _ in range(num_layers - 2):
            self.convs.append(
                GCNConv(hidden_channels, hidden_channels))
        self.convs.append(
            GCNConv(hidden_channels, out_channels))

        self.activation = F.relu

    def reset_parameters(self):
        for conv in self.convs:
            conv.reset_parameters()

    def forward(self, x, edge_index, edge_weight=None):
        for i, conv in enumerate(self.convs[:-1]):
            x = conv(x, edge_index, edge_weight)
            x = self.activation(x)
        x = self.convs[-1](x, edge_index)
        return x

def load_twitch_dataset(data_dir, train_num=3, train_ratio=0.5, valid_ratio=0.25):
    transform = T.NormalizeFeatures()
    sub_graphs = ['DE', 'PT', 'RU', 'ES', 'FR', 'EN']
    x_list, edge_index_list, y_list, env_list = [], [], [], []
    node_idx_list = []
    idx_shift = 0
    for i, g in enumerate(sub_graphs):
        torch_dataset = Twitch(root=f'{data_dir}Twitch',
                              name=g, transform=transform)
        data = torch_dataset[0]
        x, edge_index, y = data.x, data.edge_index, data.y
        x_list.append(x)
        y_list.append(y)
        edge_index_list.append(edge_index + idx_shift)
        env_list.append(torch.ones(x.size(0)) * i)
        node_idx_list.append(torch.arange(data.num_nodes) + idx_shift)

        idx_shift += data.num_nodes
    x = torch.cat(x_list, dim=0)
    y = torch.cat(y_list, dim=0)
    edge_index = torch.cat(edge_index_list, dim=1)
    env = torch.cat(env_list, dim=0)
    dataset = Data(x=x, edge_index=edge_index, y=y)
    dataset.env = env
    dataset.env_num = len(sub_graphs)
    dataset.train_env_num = train_num

    assert (train_num <= 5)

    ind_idx = torch.cat(node_idx_list[:train_num], dim=0)
    idx = torch.randperm(ind_idx.size(0))
    train_idx_ind = idx[:int(idx.size(0) * train_ratio)]
    valid_idx_ind = idx[int(idx.size(0) * train_ratio) : int(idx.size(0) * (train_ratio + valid_ratio))]
    test_idx_ind = idx[int(idx.size(0) * (train_ratio + valid_ratio)):]
    dataset.train_idx = ind_idx[train_idx_ind]
    dataset.valid_idx = ind_idx[valid_idx_ind]
    dataset.test_in_idx = ind_idx[test_idx_ind]
    dataset.test_ood_idx = [node_idx_list[-1]] if train_num>=4 else node_idx_list[train_num:]

    return dataset

# 导入数据集的时候，对数据集进行了预处理，添加了一些与环境相关的额外特征
# 定义函数 load_synthetic_dataset，
# 参数包括数据目录 data_dir，数据集名称 name，环境数量 env_num，训练环境数量 train_num，训练集比例 train_ratio，验证集比例 valid_ratio，以及一个布尔值 combine。
def load_synthetic_dataset(data_dir, name, env_num=6, train_num=3, train_ratio=0.5, valid_ratio=0.25, combine=False):
    # 这一行代码创建了一个特征归一化变换。NormalizeFeatures 是 torch_geometric.transforms 中的一个类，用于将每个节点特征向量的 L2 范数归一化为 1。具体来说，这个变换会对特征矩阵中的每一行进行归一化处理。
    transform = T.NormalizeFeatures()
    # Planetoid 类的工作机制
    # 初始化：传入数据集根目录 root、数据集名称 name 和变换 transform。
    # 下载数据：如果数据集不存在，会自动从网络下载并解压。
    # 处理数据：如果数据没有预处理，则进行预处理操作（如图结构构建、特征和标签处理等）。
    # 应用变换：在加载数据时，会应用传入的变换（如 NormalizeFeatures）。
    torch_dataset = Planetoid(root=f'{data_dir}Planetoid',
                            name=name, transform=transform)
    # 生成一个路径，用于存放预处理后的数据。
    preprocess_dir = os.path.join(data_dir, 'Planetoid', name)

    # 取得数据集的第一个图数据对象 data。
    # data可能是这样的：Data(x=[2708, 1433], edge_index=[2, 10556], y=[2708], train_mask=[2708], val_mask=[2708], test_mask=[2708])
    # x：节点特征矩阵，形状为 [2708, 1433]，表示有 2708 个节点，每个节点有 1433 个特征。
    # edge_index：边的索引矩阵，形状为 [2, 10556]，表示有 10556 条边。每条边由两个节点索引组成。
    # y：节点标签，形状为 [2708]，表示每个节点的类别标签。
    # train_mask、val_mask、test_mask：分别为训练集、验证集和测试集的掩码，形状均为 [2708]，值为 True 的位置表示该节点属于对应的数据集。
    data = torch_dataset[0]
    # 获取图的边索引 edge_index。
    edge_index = data.edge_index
    # 获取节点特征矩阵 x。
    x = data.x
    # 获取特征的维度 d。
    d = x.shape[1]

    # 更新预处理目录路径，使其指向 gen 子目录。
    preprocess_dir = os.path.join(preprocess_dir, 'gen')
    # 检查预处理目录是否存在，如果不存在，则创建它。
    if not os.path.exists(preprocess_dir):
        os.makedirs(preprocess_dir)
    # 设置生成的附加特征数量 spu_feat_num 为 10。
    spu_feat_num = 10
    # 获取数据集中的类别数量 class_num。
    class_num = data.y.max().item() + 1
    # 创建一个列表 node_idx_list，每个环境的节点索引分别存储在列表中，索引根据环境进行偏移。第1个环境{1,2,...,num_nodes}，第2个环境num_nodes+{1,2,...,num_nodes}...
    node_idx_list = [torch.arange(data.num_nodes) + i*data.num_nodes for i in range(env_num)]

    # 生成预处理文件的路径 file_path。
    file_path = preprocess_dir + f'/{class_num}-{spu_feat_num}-{env_num}.pkl'
    # 检查预处理文件是否存在，如果不存在，则进行下一步操作。
    if not os.path.exists(file_path):
        # 输出创建新合成数据的提示。
        print("creating new synthetic data...")
        # 初始化四个空列表，用于存储各环境的节点特征、边索引、标签和环境信息。
        x_list, edge_index_list, y_list, env_list = [], [], [], []
        # 初始化节点索引偏移量 idx_shift 为 0。
        idx_shift = 0

        # Generator_y = GCN_gen(in_channels=d, hidden_channels=10, out_channels=class_num, num_layers=2)
        # 初始化生成特征的 GCN 模型 Generator_x，输入通道为类别数，输出通道为生成的附加特征数量 spu_feat_num。
        Generator_x = GCN_gen(in_channels=class_num, hidden_channels=10, out_channels=spu_feat_num, num_layers=2)
        # 初始化生成噪声的线性层 Generator_noise，输入维度为环境数量 env_num，输出维度为附加特征数量 spu_feat_num。
        Generator_noise = nn.Linear(env_num, spu_feat_num)

        # 使用 torch.no_grad() 上下文管理器来禁止梯度计算（因为这是数据生成过程，不需要反向传播）。
        with torch.no_grad():
            # 遍历每个环境，生成新的特征。
            for i in range(env_num):
                # 将原始标签 data.y 转换为 one-hot 编码，并转换为浮点数。
                # numpy.squeeze(a, axis=None)
                # 其中，参数a是需要操作的数组，axis是要删除的维度。若不指定axis参数，则squeeze()方法会删除所有长度为1的维度。
                label_new = F.one_hot(data.y, class_num).squeeze(1).float()
                # 创建一个环境上下文矩阵 context_，每行表示一个节点在某个环境中的上下文信息。
                context_ = torch.zeros(x.size(0), env_num)
                context_[:, i] = 1
                # 使用 Generator_x 生成新的特征 x2，并添加噪声 Generator_noise(context_)。
                # Generator_x(label_new, edge_index)这会在图的结构基础上，根据节点的标签生成特征。
                # Generator_noise(context_)这会为每个节点生成与其环境相关的噪声特征。
                x2 = Generator_x(label_new, edge_index) + Generator_noise(context_)
                # 添加正态分布噪声，使生成特征更具随机性。
                # 这句代码是给生成的特征 x2 添加随机噪声，以增加数据的随机性和多样性。
                x2 += torch.ones_like(x2).normal_(0, 0.1)
                # 将原始特征 x 和生成的特征 x2 连接起来，形成新的特征矩阵 x_new。
                x_new = torch.cat([x, x2], dim=1)

                # 将新的特征矩阵 x_new 添加到 x_list 中。
                x_list.append(x_new)
                # 将标签 data.y 添加到 y_list 中。
                y_list.append(data.y)
                # 将偏移后的边索引 edge_index + idx_shift 添加到 edge_index_list 中。
                # 将所有节点索引平移（等于+）idx_shift
                edge_index_list.append(edge_index + idx_shift)
                # 将环境标签添加到 env_list 中，每个节点的环境标签为当前环境 i。
                # 这段代码用于处理不同环境下的节点并记录它们所属的环境
                env_list.append(torch.ones(x.size(0)) * i)
                # 更新节点索引偏移量 idx_shift，以便下一个环境的节点索引从正确的位置开始。
                idx_shift += data.num_nodes

        # 将所有环境的节点特征矩阵 x_list 连接起来，形成完整的特征矩阵 x。
        # 将 x_list 中的所有张量在第0维（行方向）连接起来，形成一个大的张量 x，表示所有环境中的节点特征。
        # 新的张量 x 的形状是 [sum(num_nodes), num_features]，其中 sum(num_nodes) 是所有环境中节点数量的总和，num_features 是节点特征的数量。
        x = torch.cat(x_list, dim=0)
        # 将所有环境的标签 y_list 连接起来，形成完整的标签 y。
        y = torch.cat(y_list, dim=0)
        # 将所有环境的边索引 edge_index_list 连接起来，形成完整的边索引 edge_index。
        edge_index = torch.cat(edge_index_list, dim=1)
        # 将所有环境的环境标签 env_list 连接起来，形成完整的环境标签 env。
        env = torch.cat(env_list, dim=0)
        # 创建一个新的数据对象 dataset，包含特征矩阵 x、边索引 edge_index 和标签 y。
        dataset = Data(x=x, edge_index=edge_index, y=y)
        # 将环境标签 env 添加到数据对象 dataset 中。
        dataset.env = env

        # 打开预处理文件路径 file_path，以二进制写入模式打开文件。
        with open(file_path, 'wb') as f:
            # 使用 pickle 序列化 dataset 对象，并写入文件。
            pkl.dump((dataset), f, pkl.HIGHEST_PROTOCOL)
    # 如果预处理文件已存在，则输出提示信息。
    else:
        print("using existing synthetic data...")
        # 打开预处理文件路径 file_path，以二进制读取模式打开文件。
        with open(file_path, 'rb') as f:
            # 用 pickle 反序列化 dataset 对象，从文件中读取数据。
            dataset = pkl.load(f)

    # 断言训练环境数量 train_num 必须小于或等于环境数量减一。
    # 合理划分训练、验证和测试环境，可以更好地评估模型的泛化能力。至少需要一个环境用于验证或测试，这样可以避免训练环境覆盖所有环境，导致无法评估模型的泛化性能。
    assert (train_num <= env_num-1)

    # 将前 train_num 个环境的节点索引连接起来，形成 ind_idx。
    ind_idx = torch.cat(node_idx_list[:train_num], dim=0)
    # 对 ind_idx 进行随机排列，生成随机索引 idx。
    idx = torch.randperm(ind_idx.size(0))
    # 根据训练比例 train_ratio 和验证比例 valid_ratio，将随机索引划分为训练集索引 train_idx_ind、验证集索引 valid_idx_ind 和测试集索引 test_idx_ind。
    train_idx_ind = idx[:int(idx.size(0) * train_ratio)]
    valid_idx_ind = idx[int(idx.size(0) * train_ratio): int(idx.size(0) * (train_ratio + valid_ratio))]
    test_idx_ind = idx[int(idx.size(0) * (train_ratio + valid_ratio)):]
    # 将训练集、验证集和测试集索引分别添加到数据对象 dataset 中。
    dataset.train_idx = ind_idx[train_idx_ind]
    dataset.valid_idx = ind_idx[valid_idx_ind]
    dataset.test_in_idx = ind_idx[test_idx_ind]

    # 如果 combine 为 True，则根据训练环境数量和环境总数确定测试集 OOD 索引 test_ood_idx，并将最后一个环境或剩余的所有环境连接起来。
    if combine:
        dataset.test_ood_idx = [node_idx_list[-1]] if train_num==env_num-1 else [torch.cat(node_idx_list[train_num:], dim=0)] # Combine three ood environments
    # 如果 combine 为 False，则直接将剩余环境的索引作为测试集 OOD 索引。
    else:
        dataset.test_ood_idx = [node_idx_list[-1]] if train_num==env_num-1 else node_idx_list[train_num:] # Test three ood environments respectively

    # 设置数据对象 dataset 的环境数量 env_num 和训练环境数量 train_num。
    dataset.env_num = env_num
    dataset.train_env_num = train_num

    # 返回数据对象 dataset。
    return dataset


def load_arxiv_dataset(data_dir, train_num=3, train_ratio=0.5, valid_ratio=0.25, inductive=True):
    from ogb.nodeproppred import NodePropPredDataset

    ogb_dataset = NodePropPredDataset(name='ogbn-arxiv', root=f'{data_dir}/ogb')

    node_years = ogb_dataset.graph['node_year']

    edge_index = torch.as_tensor(ogb_dataset.graph['edge_index'])
    node_feat = torch.as_tensor(ogb_dataset.graph['node_feat'])
    label = torch.as_tensor(ogb_dataset.labels)

    year_bound = [2005, 2010, 2012, 2014, 2016, 2018, 2021]
    env = torch.zeros(label.shape[0])
    for n in range(node_years.shape[0]):
        year = int(node_years[n])
        for i in range(len(year_bound)-1):
            if year >= year_bound[i+1]:
                continue
            else:
                env[n] = i
                break

    dataset = Data(x=node_feat, edge_index=edge_index, y=label)
    dataset.env = env
    dataset.env_num = len(year_bound)
    dataset.train_env_num = train_num

    ind_mask = (node_years < year_bound[train_num]).squeeze(1)
    idx = torch.arange(dataset.num_nodes)
    ind_idx = idx[ind_mask]
    idx_ = torch.randperm(ind_idx.size(0))
    train_idx_ind = idx_[:int(idx_.size(0) * train_ratio)]
    valid_idx_ind = idx_[int(idx_.size(0) * train_ratio): int(idx_.size(0) * (train_ratio + valid_ratio))]
    test_idx_ind = idx_[int(idx_.size(0) * (train_ratio + valid_ratio)):]
    dataset.train_idx = ind_idx[train_idx_ind]
    dataset.valid_idx = ind_idx[valid_idx_ind]
    dataset.test_in_idx = ind_idx[test_idx_ind]

    dataset.test_ood_idx = []

    for i in range(train_num, len(year_bound)-1):
        ood_mask_i = ((node_years >= year_bound[i]) * (node_years < year_bound[i+1])).squeeze(1)
        dataset.test_ood_idx.append(idx[ood_mask_i])

    return dataset


def load_elliptic_dataset(data_dir, train_num=5, train_ratio=0.5, valid_ratio=0.25):

    sub_graphs = range(0, 49)
    x_list, edge_index_list, y_list, mask_list, env_list = [], [], [], [], []
    node_idx_list = []
    idx_shift = 0
    for i in sub_graphs:
        result = pkl.load(open('{}/elliptic/{}.pkl'.format(data_dir, i), 'rb'))
        A, label, features = result
        edge_index = torch.tensor(A.nonzero(), dtype=torch.long)
        x = torch.tensor(features, dtype=torch.float)
        y = torch.tensor(label)

        x_list.append(x)
        y_list.append(y)
        mask = (y >= 0)
        edge_index_list.append(edge_index + idx_shift)
        env_list.append(torch.ones(x.size(0)) * i)
        node_idx_list.append(torch.arange(x.shape[0])[mask] + idx_shift)

        idx_shift += x.shape[0]

    x = torch.cat(x_list, dim=0)
    y = torch.cat(y_list, dim=0)
    edge_index = torch.cat(edge_index_list, dim=1)
    env = torch.cat(env_list, dim=0)
    dataset = Data(x=x, edge_index=edge_index, y=y)
    dataset.env = env
    dataset.env_num = len(sub_graphs)
    dataset.train_env_num = train_num

    ind_idx = torch.cat(node_idx_list[:train_num], dim=0)
    idx = torch.randperm(ind_idx.size(0))
    train_idx_ind = idx[:int(idx.size(0) * train_ratio)]
    valid_idx_ind = idx[int(idx.size(0) * train_ratio): int(idx.size(0) * (train_ratio + valid_ratio))]
    test_idx_ind = idx[int(idx.size(0) * (train_ratio + valid_ratio)):]
    dataset.train_idx = ind_idx[train_idx_ind]
    dataset.valid_idx = ind_idx[valid_idx_ind]
    dataset.test_in_idx = ind_idx[test_idx_ind]

    ood_margin = 4
    dataset.test_ood_idx = []
    for k in range((len(sub_graphs) - train_num*2) // ood_margin - 1):
        ood_idx_k = [node_idx_list[l] for l in range(train_num*2 + ood_margin * k, train_num*2 + ood_margin * (k + 1))]
        dataset.test_ood_idx.append(torch.cat(ood_idx_k, dim=0))
    return dataset